#include "MTCNN.h"
#include <map>
#include <iostream>



namespace onechchy
{
	const float nmsThreshold[3] = { 0.5f, 0.7f, 0.7f };
	const float meanVals[3] = { 127.5, 127.5, 127.5 };
	const float normVals[3] = { 0.0078125, 0.0078125, 0.0078125 };
	const int MIN_DET_SIZE = 12;

	MTCNN::MTCNN(const std::string& model_path)
	{
		std::vector<std::string> paramFiles = { model_path + "/det1.param", model_path + "/det2.param", model_path + "/det3.param" };
		std::vector<std::string> binFiles = { model_path + "/det1.bin", model_path + "/det2.bin", model_path + "/det3.bin" };

		int loadRet = 0;
		loadRet += pNet.load_param(paramFiles[0].c_str());
		loadRet += pNet.load_model(binFiles[0].c_str());
		loadRet += rNet.load_param(paramFiles[1].c_str());
		loadRet += rNet.load_model(binFiles[1].c_str());
		loadRet += oNet.load_param(paramFiles[2].c_str());
		loadRet += oNet.load_model(binFiles[2].c_str());

		mLoadResult = loadRet == 0;
	}
	MTCNN::MTCNN(const std::vector<std::string> paramFiles, const std::vector<std::string> binFiles)
	{
		int loadRet = 0;
		loadRet += pNet.load_param(paramFiles[0].c_str());
		loadRet += pNet.load_model(binFiles[0].c_str());
		loadRet += rNet.load_param(paramFiles[1].c_str());
		loadRet += rNet.load_model(binFiles[1].c_str());
		loadRet += oNet.load_param(paramFiles[2].c_str());
		loadRet += oNet.load_model(binFiles[2].c_str());

		mLoadResult = loadRet == 0;
	}

	MTCNN::~MTCNN()
	{
		pNet.clear();
		rNet.clear();
		oNet.clear();
	}

	bool MTCNN::loadResult()
	{
		return mLoadResult;
	}

	void MTCNN::setMinFaceSize(int minSize)
	{
		this->mMinFaceSize = minSize;
	}

	void MTCNN::detect(ncnn::Mat& img, std::vector<Bbox>& faces)
	{
		if (!this->loadResult())
		{
			return;
		}

		mImg = img;
		mImgW = img.w;
		mImgH = img.h;
		mImg.substract_mean_normalize(meanVals, normVals);

		this->PNet();
		if (this->mFirBbox.size() < 1)
		{
			return;
		}

		nms(this->mFirBbox, nmsThreshold[0]);
		refine(this->mFirBbox, mImgH, mImgW, true);

		RNet();
		if (mSecBbox.size() < 1)
		{
			return;
		}
		nms(mSecBbox, nmsThreshold[1]);
		refine(this->mSecBbox, mImgH, mImgW, true);

		ONet();
		if (mThirdBbox.size() < 1)
		{
			return;
		}
		refine(mThirdBbox, mImgH, mImgW, true);
		nms(mThirdBbox, nmsThreshold[2], "Min");
		faces = mThirdBbox;
	}

	void MTCNN::detectMaxFace(ncnn::Mat& img, std::vector<Bbox>& finalBbox)
	{
		mFirPreviousBbox.clear();
		mSecPreviousBbox.clear();
		mThirdPreviouBbox.clear();
		mFirBbox.clear();
		mSecBbox.clear();
		mThirdBbox.clear();

		this->mImg = img;
		mImgW = img.w;
		mImgH = img.h;
		this->mImg.substract_mean_normalize(meanVals, normVals);

		float minl = mImgW < mImgH ? mImgW : mImgH;
		float m = (float)MIN_DET_SIZE / mMinFaceSize;
		minl *= m;
		float factor = mPreFactor;
		std::vector<float> scales_;
		while (minl > MIN_DET_SIZE)
		{
			scales_.push_back(m);
			minl *= factor;
			m = m * factor;
		}

		std::sort(scales_.begin(), scales_.end());

		for (size_t i = 0; i < scales_.size(); ++i)
		{
			PNet(scales_[i]);
			nms(mFirBbox, nmsThreshold[0]);
			nmsTwoBoxs(mFirBbox, mFirPreviousBbox, nmsThreshold[0]);
			if (mFirBbox.size() < 1)
			{
				mFirBbox.clear();
				continue;
			}

			mFirPreviousBbox.insert(mFirPreviousBbox.end(), mFirBbox.begin(), mFirBbox.end());
			this->refine(mFirBbox, mImgH, mImgW, true);

			RNet();
			nms(mSecBbox, nmsThreshold[1]);
			nmsTwoBoxs(mSecBbox, mSecPreviousBbox, nmsThreshold[0]);
			mSecPreviousBbox.insert(mSecPreviousBbox.end(), mSecBbox.begin(), mSecBbox.end());
			if (mSecBbox.size() < 1)
			{
				mFirBbox.clear();
				mSecBbox.clear();
				continue;
			}
			this->refine(mSecBbox, mImgH, mImgW, true);

			ONet();
			if (mThirdBbox.size() < 1)
			{
				mFirBbox.clear();
				mSecBbox.clear();
				mThirdBbox.clear();
				continue;
			}
			this->refine(mThirdBbox, mImgH, mImgW, true);
			this->nms(mThirdBbox, nmsThreshold[2], "Min");

			if (mThirdBbox.size() > 0)
			{
				extractMaxFace(mThirdBbox);
				finalBbox = mThirdBbox;
				break;
			}
		}
	}

	void MTCNN::PNet()
	{
		mFirBbox.clear();
		float minl = mImgW < mImgH ? mImgW : mImgH;
		float m = (float)MIN_DET_SIZE / mMinFaceSize;
		minl *= m;
		float factor = mPreFactor;
		std::vector<float> scales;
		while (minl > MIN_DET_SIZE)
		{
 			scales.push_back(m);
			minl *= factor;
			m *= factor;
		}

		for (size_t i = 0; i < scales.size(); i++)
		{
			int ws = (int)ceil(mImgW * scales[i]);
			int hs = (int)ceil(mImgH * scales[i]);
			ncnn::Mat in;
			resize_bilinear(mImg, in, ws, hs);
			ncnn::Extractor ex = pNet.create_extractor();
			ex.set_light_mode(true);
			ex.input("data", in);
			ncnn::Mat score_, location_;
			ex.extract("prob1", score_);
			ex.extract("conv4-2", location_);
			std::vector<Bbox> boundingBox_;
			generateBbox(score_, location_, boundingBox_, scales[i]);
			nms(boundingBox_, nmsThreshold[0]);
			mFirBbox.insert(mFirBbox.end(), boundingBox_.begin(), boundingBox_.end());
			boundingBox_.clear();
		}
	}

	void MTCNN::PNet(float scale)
	{
		int hs = (int)std::ceil(mImgH * scale);
		int ws = (int)std::ceil(mImgW * scale);
		ncnn::Mat in;
		ncnn::resize_bilinear(mImg, in, ws, hs);
		ncnn::Extractor ex = pNet.create_extractor();
		ex.set_light_mode(true);
		ex.input("data", in);
		ncnn::Mat score_, location_;
		ex.extract("prob1", score_);
		ex.extract("conv4-2", location_);
		std::vector<Bbox> boundingBox_;

		generateBbox(score_, location_, boundingBox_, scale);
		nms(boundingBox_, nmsThreshold[0]);

		mFirBbox.insert(mFirBbox.end(), boundingBox_.begin(), boundingBox_.end());
		boundingBox_.clear();
	}

	void MTCNN::RNet()
	{
		this->mSecBbox.clear();
		int count = 0;
		for (std::vector<Bbox>::iterator it = mFirBbox.begin(); it != mFirBbox.end(); it++)
		{
			ncnn::Mat tempIm;
			copy_cut_border(mImg, tempIm, it->y1, mImgH - it->y2, it->x1, mImgW - it->x2);
			ncnn::Mat in;
			resize_bilinear(tempIm, in, 24, 24);
			ncnn::Extractor ex = rNet.create_extractor();
			ex.set_light_mode(true);
			ex.input("data", in);
			ncnn::Mat score_, bbox_;
			ex.extract("prob1", score_);
			ex.extract("conv5-2", bbox_);
			if ((float)score_[1] > this->mThreshold[1])
			{
				for (int channel = 0; channel < 4; channel++)
				{
					it->regreCoord[channel] = bbox_[channel];
				}
				it->area = (it->x2 - it->x1) * (it->y2 - it->y1);
				it->score = score_.channel(1)[0];
				mSecBbox.push_back(*it);
			}
		}
	}

	void MTCNN::ONet()
	{
		mThirdBbox.clear();
		for (std::vector<Bbox>::iterator it = mSecBbox.begin(); it != mSecBbox.end(); it++)
		{
			ncnn::Mat tempIm;
			copy_cut_border(mImg, tempIm, (*it).y1, mImgH - (*it).y2, (*it).x1, mImgW - (*it).x2);
			ncnn::Mat in;
			resize_bilinear(tempIm, in, 48, 48);
			ncnn::Extractor ex = oNet.create_extractor();
			ex.set_light_mode(true);
			ex.input("data", in);
			ncnn::Mat score_, bbox_, keyPoint_;
			ex.extract("prob1", score_);
			ex.extract("conv6-2", bbox_);
			ex.extract("conv6-3", keyPoint_);
			if ((float)score_[1] > mThreshold[2])
			{
				for (int channel = 0; channel < 4; channel++)
				{
					it->regreCoord[channel] = (float)bbox_[channel];
				}
				it->area = (it->x2 - it->x1) * (it->y2 - it->y1);
				it->score = score_.channel(1)[0];
				for (int num = 0; num < 5; num++)
				{
					(it->ppoint)[num] = it->x1 + (it->x2 - it->x1) * keyPoint_[num];
					(it->ppoint)[num + 5] = it->y1 + (it->y2 - it->y1) * keyPoint_[num + 5];
				}

				mThirdBbox.push_back(*it);
			}
		}
	}

	void MTCNN::generateBbox(ncnn::Mat score, ncnn::Mat location, std::vector<Bbox>& boundingBox, float scale)
	{
		const int stride = 2;
		const int cellsize = 12;

		float* p = score.channel(1);

		Bbox bbox;
		float invScale = 1.0f / scale;
		for (int row = 0; row < score.h; row++)
		{
			for (int col = 0; col < score.w; col++)
			{
				if (*p > mThreshold[0])
				{
					bbox.score = *p;
					bbox.x1 = round((stride * col + 1) * invScale);
					bbox.y1 = round((stride * row + 1) * invScale);
					bbox.x2 = round((stride * col + 1 + cellsize) * invScale);
					bbox.y2 = round((stride * row + 1 + cellsize) * invScale);
					bbox.area = (bbox.x2 - bbox.x1) * (bbox.y2 - bbox.y1);
					const int index = row * score.w + col;
					for (int channel = 0; channel < 4; channel++)
					{
						bbox.regreCoord[channel] = location.channel(channel)[index];
					}
					boundingBox.push_back(bbox);
				}
				p++;
			}
		}
	}

	void MTCNN::nms(std::vector<Bbox>& boundingBox, const float overlap_threshold, std::string modelname)
	{
		if (boundingBox.empty())
		{
			return;
		}

		std::sort(boundingBox.begin(), boundingBox.end(), [](const Bbox& a, const Bbox& b) { return a.score < b.score; });
		float IOU = 0;
		float maxX = 0;
		float maxY = 0;
		float minX = 0;
		float minY = 0;
		std::vector<int> vPick;
		int nPick = 0;
		std::multimap<float, int> vScores;
		const int num_boxes = boundingBox.size();
		vPick.resize(num_boxes);
		for (int i = 0; i < num_boxes; ++i)
		{
			vScores.insert(std::pair<float, int>(boundingBox[i].score, i));
		}

		while (vScores.size() > 0 and vPick.size() > nPick)
		{
			int last = vScores.rbegin()->second;
			vPick[nPick] = last;
			nPick += 1;
			for (std::multimap<float, int>::iterator it = vScores.begin(); it != vScores.end();)
			{
				int it_idx = it->second;
				maxX = std::max(boundingBox.at(it_idx).x1, boundingBox.at(last).x1);
				maxY = std::max(boundingBox.at(it_idx).y1, boundingBox.at(last).y1);
				minX = std::min(boundingBox.at(it_idx).x2, boundingBox.at(last).x2);
				minY = std::min(boundingBox.at(it_idx).y2, boundingBox.at(last).y2);
				maxX = ((minX - maxX + 1) > 0) ? (minX - maxX + 1) : 0;
				maxY = ((minY - maxY + 1) > 0) ? (minY - maxY + 1) : 0;

				IOU = maxX * maxY;
				if (!modelname.compare("Union"))
				{
					IOU = IOU / (boundingBox.at(it_idx).area + boundingBox.at(last).area - IOU);
				}
				else if (!modelname.compare("Min"))
				{
					IOU = IOU / ((boundingBox.at(it_idx).area < boundingBox.at(last).area) ? boundingBox.at(it_idx).area : boundingBox.at(last).area);
				}

				if (IOU > overlap_threshold)
				{
					it = vScores.erase(it);
				}
				else
				{
					it++;
				}
			}
		}

		vPick.resize(nPick);
		std::vector<Bbox> tmp_;
		tmp_.resize(nPick);
		for (int i = 0; i < nPick; i++)
		{
			tmp_[i] = boundingBox[vPick[i]];
		}
		boundingBox = tmp_;
	}

	void MTCNN::refine(std::vector<Bbox>& vecBbox, const int& height, const int& width, bool square)
	{
		if (vecBbox.empty())
		{
			std::cout << "Bbox is empty!!" << std::endl;
			return;
		}

		float bbw = 0, bbh = 0, maxSide = 0;
		float h = 0, w = 0;
		float x1 = 0, y1 = 0, x2 = 0, y2 = 0;
		for (std::vector<Bbox>::iterator it = vecBbox.begin(); it != vecBbox.end(); it++)
		{
			bbw = (*it).x2 - (*it).x1 + 1;
			bbh = (*it).y2 - (*it).y1 + 1;
			x1 = (*it).x1 + (*it).regreCoord[0] * bbw;
			y1 = (*it).y1 + (*it).regreCoord[1] * bbh;
			x2 = (*it).x2 + (*it).regreCoord[2] * bbw;
			y2 = (*it).y2 + (*it).regreCoord[3] * bbh;

			if (square)
			{
				w = x2 - x1 + 1;
				h = y2 - y1 + 1;
				maxSide = (h > w) ? h : w;
				x1 = x1 + w * 0.5 - maxSide * 0.5;
				y1 = y1 + h * 0.5 - maxSide * 0.5;
				(*it).x2 = round(x1 + maxSide - 1);
				(*it).y2 = round(y1 + maxSide - 1);
				(*it).x1 = round(x1);
				(*it).y1 = round(y1);
			}

			if ((*it).x1 < 0)
			{
				(*it).x1 = 0;
			}

			if ((*it).y1 < 0)
			{
				(*it).y1 = 0;
			}

			if ((*it).x2 > width - 1)
			{
				(*it).x2 = width - 1;
			}

			if ((*it).y2 > height - 1)
			{
				(*it).y2 = height - 1;
			}
		}
	}

	void MTCNN::extractMaxFace(std::vector<Bbox>& boundingBox)
	{
		if (boundingBox.empty())
		{
			return;
		}

		std::sort(boundingBox.begin(), boundingBox.end(), [](Bbox lsh, Bbox rsh) {
			return lsh.area > rsh.area;
			});

		for (std::vector<Bbox>::iterator itx = boundingBox.begin() + 1; itx != boundingBox.end();)
		{
			itx = boundingBox.erase(itx);
		}
	}

	void MTCNN::nmsTwoBoxs(std::vector<Bbox>& boundingBox, std::vector<Bbox>& previousBox, const float overlapThreshold, std::string modelname)
	{
		if (boundingBox.empty())
		{
			return;
		}

		std::sort(boundingBox.begin(), boundingBox.end(), [](Bbox lsh, Bbox rsh) {
			return lsh.score < rsh.score;
			});
		float IOU = 0;
		float maxX = 0;
		float maxY = 0;
		float minX = 0;
		float minY = 0;
		for (std::vector<Bbox>::iterator ity = previousBox.begin(); ity != previousBox.end(); ity++)
		{
			for (std::vector<Bbox>::iterator itx = boundingBox.begin(); itx != boundingBox.end();)
			{
				int i = itx - boundingBox.begin();
				int j = ity - boundingBox.begin();
				maxX = std::max(boundingBox.at(i).x1, previousBox.at(j).x1);
				maxY = std::max(boundingBox.at(i).y1, previousBox.at(j).y1);
				minX = std::min(boundingBox.at(i).x2, previousBox.at(j).x2);
				minY = std::min(boundingBox.at(i).y2, previousBox.at(j).y2);

				maxX = ((minX - maxX + 1) > 0) ? (minX - maxX + 1) : 0;
				maxY = ((minY - maxY + 1) > 0) ? (minY - maxY + 1) : 0;

				IOU = maxX * maxY;
				if (!modelname.compare("Union"))
				{
					IOU = IOU / (boundingBox.at(i).area + previousBox.at(j).area - IOU);
				}
				else if (!modelname.compare("Min"))
				{
					IOU = IOU / ((boundingBox.at(i).area < previousBox.at(j).area) ? boundingBox.at(i).area : previousBox.at(j).area);
				}

				if (IOU > overlapThreshold && boundingBox.at(i).score > previousBox.at(j).score)
				{
					itx = boundingBox.erase(itx);
				}
				else
				{
					itx++;
				}
			}
		}
	}

}